import { BaseRetriever } from "@langchain/core/retrievers";
import { OllamaEmbeddings } from "@langchain/ollama";
import { MilvusClient } from '@zilliz/milvus2-sdk-node';
import { Milvus } from '@langchain/community/vectorstores/milvus'

import { CallbackManagerForRetrieverRun } from "@langchain/core/callbacks/manager";

class CustomRetriever extends BaseRetriever {

    lc_namespace = ["langchain", "retrievers", "custom_retriever"];


    private embeddingModel: OllamaEmbeddings;

    private milvus: MilvusClient;

    private Milvus: Milvus;

    constructor() {
        super();
        this.embeddingModel = new OllamaEmbeddings({ // 初始化 embedding 模型，用于向量化文本，用于后续的检索。

            baseUrl: 'http://127.0.0.1:11434', // ollama服务的地址

            model: 'nomic-embed-text:latest', // 通过ollama部署的模型名称
        })

        this.milvus = new MilvusClient({
            address: 'http://127.0.0.1:19530',
        });
        this.Milvus = new Milvus(this.embeddingModel, { 
            clientConfig: {
                address: 'http://127.0.0.1:19530',
            }
        })
    }

    static lc_name() {
        return "CustomRetriever";
    }

    _getType() {
        return 'custom_retriever';
    }

    async _getRelevantDocuments(
        query,
        _callbacks?: CallbackManagerForRetrieverRun
    ): Promise<any[]> {
        console.log(42, '_getRelevantDocuments')
        // 自定义检索逻辑
        const relevantDocs = this.searchRelevantDocs(query);

        return relevantDocs as any;
    }

    async aGetRelevantDocuments (
        query,
        _callbacks?: CallbackManagerForRetrieverRun
    ) {
        const relevantDocs = this.searchRelevantDocs(query);

        return relevantDocs as any;
    }
    async searchRelevantDocs(question) {
        // 实现具体的文档检索逻辑
        // 这里可以调用你的检索算法或数据库查询
        // 示例：从一个简单的数组中检索相关文档

        // 1. 问题向量化
        const vector = await this.embeddingModel.embedQuery(question);
        console.log('vector', vector);

        // 2. 向量检索，返回相关文档
        const relevantDocs = await this.milvus.search({
            collection_name: 'xiaoming', // 替换为你的集合名称
            data: vector,
            output_fields: ['document', 'id'], // 替换为你的字段名称
        })
        console.log('relevantDocs', relevantDocs);

        return relevantDocs;
    }
}

export default CustomRetriever;